// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{ffi::OsStr, path::Path, str::FromStr};

use fixedbitset::FixedBitSet;
use indicatif::MultiProgress;
use log::{error, info, trace};
use petgraph::{graph::NodeIndex, visit::EdgeRef};
use tokio::process::Command;

use crate::{build::{envs::{prepare_env_vars_for_deploy, prepare_env_vars_for_deploy_in_rollback}, BuildContext}, config::system::{toolchain::{ToolchainConfig, ToolchainTripleEnvsConfig}, SystemConfig}, database::{deploy_stat::DeployStats, local::{pkg_stat_path, PKGS_CONFIG_DIR}, mirror::Mirrors, pkg_stat::{PkgStatFile, StatSpec}}, graph::{build_dep_graph::{DepGraph, DepNode}, DepKind}, Arch};

use super::{fhs::{build_fhs, FhsBuildError}, merge::merge_pkg, MergeError, MergeMode};

#[derive(Debug, thiserror::Error)]
pub enum DeployError {
    #[error("failed to build fhs")]
    FhsBuildFailed(#[from] FhsBuildError),
    #[error("failed to merge package")]
    MergeFailed(#[from] MergeError),
    #[error("failed to run deploy script")]
    ScriptFailed,
    #[error(transparent)]
    IoError(#[from] tokio::io::Error),
}

pub fn collect_deploy_deps(dep_graph: &DepGraph, ids: Vec<NodeIndex>) -> FixedBitSet {
    let mut to_deploy = dep_graph.visit_map();
    let mut unresolved = ids;
    while let Some(id) = unresolved.pop() {
        if to_deploy.contains(id.index()) {
            continue;
        }

        to_deploy.set(id.index(), true);
        for dep_edge in dep_graph.dep_edges(id) {
            let dep_node_id = dep_edge.source();
            match dep_edge.weight().kind {
                DepKind::Build | DepKind::HostBuild => {},
                DepKind::Runtime => {
                    let config = unsafe { dep_edge.weight().get_config() };
                    if config.is_none_or(|c| c.recursive.is_none_or(|r| r)) {
                        unresolved.push(dep_node_id)
                    }
                },
            }
        }
    }

    to_deploy
}

pub async fn install_deploy_runtime_deps(
    to_deploy: &FixedBitSet,
    spec_strings: &[String],
    dep_graph: DepGraph,
    sorted_ids: Vec<NodeIndex>,
    toolchain_config: ToolchainConfig,
    system_config: &SystemConfig,
    multi_progress: MultiProgress,
) -> anyhow::Result<(DepGraph, ToolchainConfig)> {
    let mut need_install = false;
    for id in to_deploy.ones() {
        if !pkg_stat_path(&spec_strings[id]).exists() {
            trace!("{} not installed", &spec_strings[id]);
            need_install = true;
        }
    }
    if need_install {
        info!("Installing the packages to be deployed...");
        let mirror = Mirrors::fetch_index_configs(system_config.default_mirror_urls().map(|s| s.as_str())).await;
        let mut build_context = BuildContext::new(
            dep_graph,
            sorted_ids,
            toolchain_config,
            mirror,
            false,
        ).await?;
        build_context.build_all(multi_progress).await?;
        Ok(build_context.unwrap_arc())
    } else {
        Ok((dep_graph, toolchain_config))
    }
}

pub async fn filter_deployed_pkgs(
    mut to_deploy: FixedBitSet,
    spec_strings: &[String],
    root: &Path,
)-> Result<FixedBitSet, anyhow::Error> {
    let stats_dir = DeployStats::deploy_stat_dir(root);
    if let Some((current_stats, _)) = DeployStats::load_current(&stats_dir).await? {
        let deployed_pkgs = current_stats.deployed_pkgs();
        for id in to_deploy.clone().ones() {
            if !pkg_stat_path(&spec_strings[id]).exists() {
                return Err(anyhow::anyhow!("{} not installed automatically, please install the pkg manually", &spec_strings[id]));
            }
            if deployed_pkgs.contains(&spec_strings[id].as_str()) {
                to_deploy.set(id, false);
            }
        }
    }

    Ok(to_deploy)
}

pub fn collect_deploy_deps_for_build(dep_graph: &DepGraph, start_id: NodeIndex) -> FixedBitSet {
    let mut to_deploy = dep_graph.visit_map();
    let mut build_deploy = dep_graph.visit_map();
    let mut runtime_deploy = dep_graph.visit_map();

    let mut build_unresolved = vec![];
    let mut unresolved = vec![];

    let node = dep_graph.node(start_id).unwrap();
    for kind in node.config.toolchains.iter() {
        let ranges = dep_graph.id_map.subtoolchain_ranges.get(kind).unwrap();
        unresolved.extend(ranges.host_toolchain_range.clone().map(NodeIndex::new));

        let need_cross = !node.spec.arch.is_host()
            && !node.config.build_options()
                .is_some_and(|o| o.dont_need_cross_compile);
        if need_cross && !dep_graph.id_map.temp_toolchain_range.contains(&start_id.index()) {
            if let Some(range) = ranges.cross_toolchain_ranges.get(&node.spec.arch) {
                if !range.contains(&start_id.index()) {
                    unresolved.extend(range.clone().map(NodeIndex::new));
                }
            }
        };
    }

    for dep_edge in dep_graph.dep_edges(start_id) {
        let dep_node_id = dep_edge.source();
        match dep_edge.weight().kind {
            DepKind::Build => build_unresolved.push(dep_node_id),
            DepKind::HostBuild => unresolved.push(dep_node_id),
            DepKind::Runtime => {},
        }
    }

    while let Some(id) = unresolved.pop() {
        if runtime_deploy.contains(id.index()) {
            continue;
        }
        runtime_deploy.set(id.index(), true);
        to_deploy.set(id.index(), true);
        
        for dep_edge in dep_graph.dep_edges(id) {
            let dep_node_id = dep_edge.source();
            match dep_edge.weight().kind {
                DepKind::Build | DepKind::HostBuild => {},
                DepKind::Runtime => {
                    let config = unsafe { dep_edge.weight().get_config() };
                    if config.is_none_or(|c| c.recursive.is_none_or(|r| r)) {
                        unresolved.push(dep_node_id);
                    }
                },
            }
        }
    }

    while let Some(id) = build_unresolved.pop() {
        if build_deploy.contains(id.index()) {
            continue;
        }
        build_deploy.set(id.index(), true);
        to_deploy.set(id.index(), true);

        for dep_edge in dep_graph.dep_edges(id) {
            let dep_node_id = dep_edge.source();
            match dep_edge.weight().kind {
                DepKind::Build => {
                    let config = unsafe { dep_edge.weight().get_config() };
                    if config.is_some_and(|c| c.recursive.is_some_and(|r| r)) {
                        build_unresolved.push(dep_node_id);
                    }
                },
                DepKind::HostBuild | DepKind::Runtime => {},
            }
        }
    }

    to_deploy
}

pub async fn deploy_pkgs<I>(
    root: &Path,
    dep_graph: &DepGraph,
    spec_strings: &[String],
    toolchain_triple_envs: &ToolchainTripleEnvsConfig,
    extra_envs: &[(&OsStr, &OsStr)],
    pkgs: I,
    arch: Arch,
    impure: bool,
    mode: MergeMode,
    single_arch: bool,
) -> Result<(), DeployError>
where
    I: Iterator<Item = NodeIndex>
{
    build_fhs(root).await?;

    for id in pkgs {
        let spec_str = if impure {
            &format!("{}-impure", spec_strings[id.index()])
        } else {
            &spec_strings[id.index()]
        };
        let node = dep_graph.node(id).unwrap();
        if !single_arch || node.spec.arch == arch {
            let node = dep_graph.node(id).unwrap();
            merge_pkg(root, node.config.deploy.as_ref(), node.spec.arch, spec_str, arch, mode)?;
            if let Some(c) = node.config.deploy.as_ref() {
                if let Some(script) = c.script.as_ref() {
                    let command = deploy_script_cmd(root, script, node, arch, toolchain_triple_envs, extra_envs.iter());
                    run_deploy_script(command).await?;
                }
            }
        }
    }
    Ok(())
}

pub async fn deploy_pkgs_for_rollback<'a, I>(
    root: &Path,
    to_deploy_pkgs: I,
    arch: &Arch,
) -> anyhow::Result<()>
where I: Iterator<Item = (&'a str, MergeMode)>
{
    build_fhs(root).await?;

    for (pkg, mode) in to_deploy_pkgs {
        let stat_spec = StatSpec::from_str(pkg)?;
        let stat_path = pkg_stat_path(pkg);
        let mut stat_file = PkgStatFile::open(&stat_path).await?;
        let features = stat_file.read_features().await?;
        let deploy_config = stat_file.read_deploy_config().await?;
        merge_pkg(root, deploy_config.as_ref(), stat_spec.arch, pkg, *arch, mode)?;
        if let Some(c) = deploy_config.as_ref() {
            if let Some(script) = c.script.as_ref() {
                let command = deploy_script_cmd_for_rollback(
                    root,
                    script,
                    &stat_spec,
                    &features,
                    *arch,
                );
                run_deploy_script(command).await?;
            }
        }
    }

    Ok(())
}
async fn run_deploy_script(mut command: Command) -> Result<(), DeployError> {
    let output = command
        .kill_on_drop(true)
        .output()
        .await?;
    if !output.status.success() {
        error!(
            "Deploy error: command exit with {}\ncommand: {:?}\nstdout:\n{}\nstderr:\n{}",
            output.status,
            command,
            &String::from_utf8_lossy(&output.stdout),
            &String::from_utf8_lossy(&output.stderr),
        );
        return Err(DeployError::ScriptFailed);
    }
    Ok(())
}

fn deploy_script_cmd<'a, I>(
    root: &Path,
    script: &String,
    node: &DepNode,
    arch: Arch,
    toolchain_triple_envs: &ToolchainTripleEnvsConfig,
    extra_envs: I,
) -> Command
where I: Iterator<Item = &'a (&'a OsStr, &'a OsStr)>
{
    let script_path = Path::new(PKGS_CONFIG_DIR.as_str())
        .join(node.spec.pkg.to_string())
        .join(script);
    let mut cmd = Command::new(script_path);
    cmd.current_dir(root);
    prepare_env_vars_for_deploy(node, arch, toolchain_triple_envs, extra_envs, |k, v| {
        cmd.env(k, v);
    });
    cmd
}

fn deploy_script_cmd_for_rollback(
    root: &Path,
    script: &String,
    spec: &StatSpec,
    features: &[String],
    arch: Arch,
) -> Command {
    let script_path = Path::new(PKGS_CONFIG_DIR.as_str())
        .join(spec.pkg.to_string())
        .join(script);
    let mut cmd = Command::new(script_path);
    cmd.current_dir(root);
    prepare_env_vars_for_deploy_in_rollback(spec, features, arch, |k, v| {
        cmd.env(k, v);
    });
    cmd
}